#include "horizon/hbmodel_manager.h"
#include "base_util/utils.h"
#include "base_util/math_compute.h"

void performance_test(ai::HBModelManager* model, std::string yuv_path, cv::Size yuv_size, std::string model_id, int max_channel_size)
{
  int ret = 0;
  stream::ImageBlob* blob = new stream::ImageBlob(stream::ImageBlobMode_YUV);

  std::vector<char> buffer; 
  std::ifstream infile;
  infile.open(yuv_path);
  infile.seekg(0,std::ios::end);//跳转到文件末尾
  long len = infile.tellg();//获取偏移量-文件长度
  infile.seekg(0,std::ios::beg);//跳转到文件起始位置
  buffer.resize(len);//申请空间
  infile.read(buffer.data(),len);//读取文件数据
  infile.close();//关闭文件

  std::vector<int> fps_cnt(max_channel_size, 0);
  std::vector<float> npu_infer_fps(max_channel_size, 0);
  int64_t start_time = 0, end_time = 0;
  start_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();


  std::vector<int> infer_ids;
  // 模拟通道号
  int channel_id = 0;
  while(1)
  {
    
    // 喂数据
    int infer_id = -1;
    blob->yuv = buffer;
    blob->ori_im_shape = {yuv_size.width, yuv_size.height};
    blob->source_id = channel_id;
    blob->model_id=model_id;
    blob->ori_im_shape={640,640};    
    model->set_data(blob->model_id, infer_id, *blob);
    // 数据设置成功 此数据必然会被拿去推理
    if (infer_id >= 0) {
      channel_id++;
      if (channel_id >= max_channel_size) {channel_id = 0;}
      infer_ids.push_back(infer_id);
    }

    // 拿数据
    {
      infer_id = infer_ids.front();
      ai::InferResult infer_res;
      ret = model->get_data(model_id, infer_id, infer_res);
      if (ret == 0) {
        fps_cnt[infer_res.source_id]++;
        for (auto& det_info : infer_res.infer_res) {delete det_info;}
      }
    }

    end_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
    if (end_time - start_time > 5 * 1e3) {
      for (int idx = 0; idx< max_channel_size; idx++) {
        npu_infer_fps[idx] = 1.0 / ((end_time - start_time) / (fps_cnt[idx] * 1e3));
        std::cout << "channel:"<< idx << " infer_fps:" << npu_infer_fps[idx] << std::endl;
        fps_cnt[idx] = 0;
      }
      start_time = end_time;
    }

    std::this_thread::sleep_for(milliseconds(3));

  }



}


int infer_yuv(ai::HBModelManager* model, std::vector<ai::InitModelData>& init_model_datas, std::string yuv_path, cv::Size yuv_size) {
  int ret = 0;

  std::vector<char> buffer; 
  std::ifstream infile;
  infile.open(yuv_path);
  infile.seekg(0,std::ios::end);//跳转到文件末尾
  long len = infile.tellg();//获取偏移量-文件长度
  infile.seekg(0,std::ios::beg);//跳转到文件起始位置
  buffer.resize(len);//申请空间
  infile.read(buffer.data(),len);//读取文件数据
  infile.close();//关闭文件

  static int frame_cnt = 0;
  stream::ImageBlob blob(stream::ImageBlobMode_YUV); 
  blob.source_id = frame_cnt++;
  blob.model_id = init_model_datas[0].model_id;
  blob.yuv = buffer;
  blob.yuv_size = buffer.size();
  blob.ori_im_shape = {yuv_size.height, yuv_size.width};
    

  int rec_idx = 0;
  model->set_data(blob.model_id, rec_idx, blob);
  if (rec_idx < 0) {
    std::cout << "rec_idx: " << rec_idx << std::endl;
    return -1;
  }

  ai::InferResult infer_res;
  while ((ret = model->get_data(blob.model_id, rec_idx, infer_res)) != 0) {
    std::this_thread::sleep_for(milliseconds(3));
  }
  std::cout << "det_infos.size():" << infer_res.infer_res.size() << std::endl;

  cv::Mat img;
  nv122bgr(blob.yuv, yuv_size, img);
  for (auto& det_info : infer_res.infer_res) {
    ai::DetInfo* di = (ai::DetInfo*)det_info;
    cv::rectangle(img, di->det_box, cv::Scalar(0, 0, 255), 2);
    cv::putText(img, di->category + ":" + std::to_string(di->score), cv::Point(di->det_box.x, di->det_box.y - 20), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 0, 255), 2);
    delete di;
  }
  cv::imwrite(yuv_path+".jpg", img);
  
  return 0;
}


int infer_img(ai::HBModelManager* model, std::vector<ai::InitModelData>& init_model_datas, std::string img_dir) {
  int ret = 0;

  std::vector<boost::filesystem::path> all_files;
  get_files_in_dir(img_dir, all_files, {".jpg"});
  std::vector<std::string> all_img_paths;
  for (auto& img_path : all_files) { all_img_paths.push_back(img_path.string());}
  std::sort(all_img_paths.begin(), all_img_paths.end(), [](std::string& a, std::string& b){return a < b;});

  int frame_cnt = 0;
  stream::ImageBlob blob(stream::ImageBlobMode_YUV); 
  blob.model_id = init_model_datas[0].model_id;

  for (auto& img_path : all_img_paths) {
    cv::Mat bgr_img = cv::imread(img_path);
    cv::Mat nv12_img;
    bgr2nv12(bgr_img, nv12_img);

    std::vector<char> buffer; 
    buffer.resize(nv12_img.cols * nv12_img.rows);
    memcpy(buffer.data(), nv12_img.data, nv12_img.cols * nv12_img.rows);

    blob.source_id = frame_cnt++;
    blob.yuv = buffer;
    blob.yuv_size = buffer.size();
    blob.ori_im_shape = {bgr_img.rows, bgr_img.cols};
      

    int rec_idx = 0;
    model->set_data(blob.model_id, rec_idx, blob);
    if (rec_idx < 0) {
      std::cout << "rec_idx: " << rec_idx << std::endl;
      return -1;
    }

    ai::InferResult infer_res;
    while ((ret = model->get_data(blob.model_id, rec_idx, infer_res)) != 0) {
      std::this_thread::sleep_for(milliseconds(3));
    }
    std::cout << "det_infos.size():" << infer_res.infer_res.size() << std::endl;

    for (auto& det_info : infer_res.infer_res) {
      ai::DetInfo* di = (ai::DetInfo*)det_info;
      cv::rectangle(bgr_img, di->det_box, cv::Scalar(0, 0, 255), 2);
      cv::putText(bgr_img, di->category + ":" + std::to_string(di->score), cv::Point(di->det_box.x, di->det_box.y - 20), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 0, 255), 2);
      delete di;
    }
    cv::imwrite(img_path+".draw.jpg", bgr_img);
  }

  return 0;
}

int infer_cls_video(ai::HBModelManager* model, std::vector<ai::InitModelData>& init_model_datas, std::string img_dir) {
  int ret = 0;

  std::vector<boost::filesystem::path> all_files;
  get_files_in_dir(img_dir, all_files, {".jpg"});
  std::vector<std::string> all_img_paths;
  for (auto& img_path : all_files) { all_img_paths.push_back(img_path.string());}
  std::sort(all_img_paths.begin(), all_img_paths.end(), [](std::string& a, std::string& b){return a < b;});

  int frame_cnt = 0;
  int infer_id = 0;
  int plate_rec_idx = 1;
  stream::ImageBlob video_blob(stream::ImageBlobMode_NONE);
  // 填充8张图片
  for (int idx = 0; idx < 8 && idx < all_img_paths.size(); idx++) {
    cv::Mat img = cv::imread(all_img_paths[idx]);
    video_blob.imgs.push_back(img);
  }

  video_blob.source_id = frame_cnt;
  video_blob.model_id = init_model_datas[0].model_id;
  model->set_data(video_blob.model_id, infer_id, video_blob);
  std::cout << "infer_id:" << infer_id << std::endl;
  if (infer_id < 0) {return 0;}

  ai::InferResult infer_res;
  while ((ret = model->get_data(video_blob.model_id, infer_id, infer_res)) != 0) {
    std::this_thread::sleep_for(milliseconds(3));
  }
  std::cout << "det_infos.size():" << infer_res.infer_res.size() << std::endl;

  for (auto& cls_info : infer_res.infer_res) {
    ai::ClsInfo* cls = (ai::ClsInfo*)cls_info;
    for (int cls_idx = 0; cls_idx < cls->top5_score.size(); cls_idx++) {
      std::cout << "class:" << cls->top5_class[cls_idx] << ", score:" << cls->top5_score[cls_idx] << std::endl;
    }
    delete cls;
  }
  frame_cnt++;
}

int main(int argc, char *argv[]){
  std::string model_dir = argv[1];
  std::string yuv_path = argv[2];
  int yuv_w = std::stoi(argv[3]);
  int yuv_h = std::stoi(argv[4]);
  cv::Size yuv_size = cv::Size(yuv_w, yuv_h); 

  int ret = 0;
  LogInfo *log_info = new LogInfo(3, "log", "log");
  ai::HBModelManager* model = new ai::HBModelManager(log_info);

  // // 测试加密模型时需要
  // //初始化授权
  // ret = model->init_param(log_info);
  // if(ret != 0){
  //   std::cout << "ERROR. init_param ret:" << ret << std::endl; 
  //   return ret;
  // }
  // //获取所支持的算法标签
  // std::vector<au::authModelLabel> auth_mdls;
  // ret = model->get_all_model_info(auth_mdls);


  std::vector<ai::InitModelData> init_model_datas;
  // 原始
  // init_model_datas.push_back(ai::InitModelData(model_dir,"det/base.9", "ssdh67fgw", "base.9", ai::model_det, 0, "cfg.yaml"));
  // init_model_datas.push_back(ai::InitModelData(model_dir,"det/common5", "7fgwssdh6", "common5", ai::model_det, 0, "cfg.yaml"));
  init_model_datas.push_back(ai::InitModelData(model_dir,"cls/firesmoke", "ssdh67fgw", "firesmoke", ai::model_cls_video, 1, "cfg.yaml"));
  // 加密
  // init_model_datas.push_back(ai::InitModelData(model_dir, "", "det_chongdianzhan_yolom.20250410_001", "common5", ai::model_det, 0, "det_chongdianzhan_yolom.20250410_001.yaml"));


  ret = model->init_model(init_model_datas, 0);
  if (ret != 0) { 
    std::cout << "init_model ret: " << ret << std::endl;
    return ret; 
  }

  std::this_thread::sleep_for(std::chrono::milliseconds(3000));
  
  // // 性能测试
  int max_size = 2;
  // performance_test(model, yuv_path, yuv_size, init_model_datas[0].model_id, max_size);

  // yuv 识别
  // infer_yuv(model, init_model_datas, yuv_path, yuv_size);
  // img 识别
  // infer_img(model, init_model_datas, yuv_path);
  // 视频分类
  infer_cls_video(model, init_model_datas, yuv_path);

  return 0;
}



/*
  cv::Mat bgr_img = cv::imread(yuv_path);
  cv::Mat yuv_img;
  bgr2nv12(bgr_img, yuv_img);

  int data_size = bgr_img.cols * bgr_img.rows * 3/2;
  std::ofstream imageyuv(yuv_path+".yuv", std::ios::binary);
  if(imageyuv.is_open()) {
    imageyuv.write((char*)yuv_img.data, data_size);
    imageyuv.close();
  }
  return 0;
*/